# %%
import numpy as np
import matplotlib.pyplot as plt
from sklearn.neural_network import MLPRegressor
from sklearn.metrics import mean_squared_error


# %%
# 洛伦兹方程的函数定义
def lorenz(x, y, z, sigma=10, r=28, b=8/3):
    dx = -sigma * x + sigma * y
    dy = -x * z + r * x - y
    dz = x * y - b * z
    return dx, dy, dz

# 使用4阶Runge-Kutta方法生成洛伦兹吸引子的时间序列数据
def generate_lorenz_data(timesteps, dt=0.01):
    x, y, z = 1.0, 1.0, 1.0  # 初始条件
    data = []
    for _ in range(timesteps):
        dx1, dy1, dz1 = lorenz(x, y, z)
        x += (dx1 ) * dt 
        y += (dy1) * dt 
        z += (dz1 ) * dt 
        data.append([x, y, z])
    return np.array(data)

# 生成数据
timesteps = 3000
data = generate_lorenz_data(timesteps)




# %%
fig = plt.figure(figsize=(10, 8))
ax = fig.add_subplot(111, projection='3d')

ax.plot(data[:, 0], data[:, 1], data[:, 2], lw=0.5, color='blue')
ax.set_xlabel('x')
ax.set_ylabel('y')
ax.set_zlabel('z')
ax.set_title('3D Lorenz Attractor')
plt.show()

# %%
# 构造输入特征，每个样本是过去20个时间点的x,y,z，用0填充不足的部分
data_extended = np.vstack([np.zeros((19, 3)), data])  # 前面填充19行0
X_new = np.zeros((len(data), 20 * 3))
for k in range(len(data)):
    window = data_extended[k:k+20]  # 获取20个时间点的数据
    X_new[k] = window.flatten()

# 划分训练集和测试集
X_train = X_new[:700]
y_train = data[:700, 0]
X_test = X_new[700:]
y_test = data[700:, 0]

# %%
# 定义多层感知机回归模型
mlp = MLPRegressor(hidden_layer_sizes=(200), max_iter=100, learning_rate_init=0.1, alpha=0.0001, solver='adam',activation = 'relu')

# 训练模型
mlp.fit(X_train, y_train)

# %%
# 绘制均方误差与epoch的关系
train_errors = []
for i in range(1, 51):
    mlp.max_iter = i
    mlp.fit(X_train, y_train)
    train_errors.append(mean_squared_error(y_train, mlp.predict(X_train)))

plt.plot(range(1, 51), train_errors, label="Training Error")
plt.xlabel("Epochs")
plt.ylabel("Mean Squared Error")
plt.title("Training MSE vs Epochs")
plt.show()

# 使用训练好的模型进行预测
y_pred = mlp.predict(X_test)

# 绘制预测结果与真实数据的对比
plt.plot(y_test, label="True Values")
plt.plot(y_pred, label="Predicted Values")
plt.xlabel("Time Step")
plt.ylabel("Z Value")
plt.title("Prediction vs Actual Lorenz Attractor")
plt.legend()
plt.show()


